# -*- coding: UTF-8 -*-
"""
@Date    ：2025/9/28 16:26 
@Author  ：Liu Yuezhao
@Project ：bert 
@File    ：event_target_dataset.py
@IDE     ：PyCharm 
"""
from torch.utils.data import Dataset
import pandas as pd
import numpy as np
import ast
import torch

class MBDDataset(Dataset):
    def __init__(self,
                 trx_target_data: pd.DataFrame,
                 trx_dict: dict,
                 max_seq_len: int,
                 device,
                 pad_token_id=0,
                 cols_used=['event_subtype'],
                 target_used=['bcard_target'],
                 time_col='event_time',
                 time_unit='hour'):
        """
        用于E2E和fine tune的时间信息事件序列 Dataset
        Args:
            trx_target_data: DataFrame，每行是一个用户/序列
            trx_dict: 词表字典，包含各列的 vocab_size, cls_token_id 等
            max_seq_len: 最大序列长度（含 CLS/SEP）
            device: 数据加载到的设备（如 'cuda' 或 'cpu'）
            pad_token_id: 填充 token ID
            cols_used: 用于 MLM 的事件列名列表，如 ['event_subtype']
            target_used: 作为labels的列名
            time_col: 时间戳列名
            time_unit: 时间单位，'day' 或 'hour' 或 'minute' 或 'second'
        """
        super().__init__()
        self.data = trx_target_data.reset_index(drop=True)
        self.trx_dict = trx_dict
        self.max_seq_len = max_seq_len
        self.device = device
        self.pad_token_id = pad_token_id
        self.cols_used = cols_used
        self.target_used = target_used
        self.time_col = time_col
        self.time_unit = time_unit

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        return self._process_seq(self.data.iloc[index].to_dict())

    def _parse_timestamps(self, timestamps):
        """解析时间戳"""
        if isinstance(timestamps, str):
            try:
                timestamps = ast.literal_eval(timestamps)
            except Exception as e:
                raise ValueError("Cannot parse string as list of timestamps")

        ts_array = np.array(timestamps)
        if np.issubdtype(ts_array.dtype, np.number):
            # 智能判断时间戳单位（秒、毫秒、微秒）
            non_zero = ts_array[ts_array > 0]
            if len(non_zero) == 0:
                raise ValueError("All timestamps are zero or negative.")
            avg_val = float(np.mean(non_zero))

            if avg_val < 1e10:      # 秒级
                scale = 1e9
            elif avg_val < 1e13:    # 毫秒级
                scale = 1e6
            elif avg_val < 1e16:    # 微秒级
                scale = 1e3
            else:
                raise ValueError(f"Unrealistic timestamp magnitude: {avg_val}")

            return (ts_array * scale).astype('datetime64[ns]')
        else:
            return pd.to_datetime(ts_array).values

    def _compute_age_in_days(self, timestamps):
        """计算每个事件距离序列第一个事件的时间差"""
        if len(timestamps) == 0:
            return np.array([])
        first_time = timestamps[0]
        deltas = timestamps - first_time
        if self.time_unit == 'day':
            return deltas.astype('timedelta64[D]').astype(float)
        elif self.time_unit == 'hour':
            return deltas.astype('timedelta64[h]').astype(float)
        elif self.time_unit == 'minute':
            return deltas.astype('timedelta64[m]').astype(float)
        elif self.time_unit == 'second':
            return deltas.astype('timedelta64[s]').astype(float)
        else:
            raise ValueError("Unrecognized time unit.")

    def _compute_interval_in_days(self, timestamps):
        """计算相邻事件之间的时间间隔"""
        if len(timestamps) <= 1:
            return np.array([0.0])
        deltas = np.diff(timestamps)
        if self.time_unit == 'day':
            interval = deltas.astype('timedelta64[D]').astype(float)
        elif self.time_unit == 'hour':
            interval = deltas.astype('timedelta64[h]').astype(float)
        elif self.time_unit == 'minute':
            interval = deltas.astype('timedelta64[m]').astype(float)
        elif self.time_unit == 'second':
            interval = deltas.astype('timedelta64[s]').astype(float)
        else:
            raise ValueError("Unrecognized time unit.")
        return np.concatenate([[0.0], interval])

    def _compute_same_event_interval_in_days(self, timestamps, event_types):
        """计算每个事件距离上一次同类事件的时间间隔"""
        if len(timestamps) == 0:
            return np.array([])

        interval = np.zeros(len(timestamps))
        last_occurrence = {}

        for i in range(len(timestamps)):
            event_type = event_types[i]
            current_time = timestamps[i]

            if event_type in last_occurrence:
                prev_time = last_occurrence[event_type]
                delta = current_time - prev_time
                if self.time_unit == 'day':
                    interval[i] = delta.astype('timedelta64[D]').astype(float)
                elif self.time_unit == 'hour':
                    interval[i] = delta.astype('timedelta64[h]').astype(float)
                elif self.time_unit == 'minute':
                    interval[i] = delta.astype('timedelta64[m]').astype(float)
                elif self.time_unit == 'second':
                    interval[i] = delta.astype('timedelta64[s]').astype(float)

            last_occurrence[event_type] = current_time

        return interval

    def _pad_time_sequence(self, seq):
        """
        专门用于时间特征的填充：
        - 截断保留最近 max_len-2 个事件
        - [CLS] 和 [SEP] 位置的时间值设为 0.0
        - 左侧用 0.0 填充
        """
        max_len_wo_special = self.max_seq_len - 2
        truncated = seq[-max_len_wo_special:] if len(seq) > max_len_wo_special else seq

        padded = [0.0] * (self.max_seq_len - len(truncated) - 2)  # 左侧填充 0
        padded += [0.0]  # [CLS] 位置时间特征为 0
        padded += truncated
        padded += [0.0]  # [SEP] 位置时间特征为 0
        return padded

    def _add_special_token(self, seq, col):
        """
        逻辑修复：先对原始 seq 左补 pad 到 max_seq_len - 2，再在首尾加 [CLS] 和 [SEP]
        最终格式: [CLS] + [PAD, ..., PAD, truncated_seq] + [SEP]
        """
        if isinstance(seq, np.ndarray):
            seq = seq.tolist()

        max_len_wo_special = self.max_seq_len - 2  # 留出 [CLS] 和 [SEP]

        # 1. 截断原始序列（保留右侧，即最近的行为）
        truncated = seq[-max_len_wo_special:] if len(seq) > max_len_wo_special else seq

        # 2. 对 truncated 左补 pad 到 max_len_wo_special
        padding_length = max_len_wo_special - len(truncated)
        padded_seq = [self.pad_token_id] * padding_length + truncated  # [PAD, ..., token1, token2, ...]

        # 3. 添加特殊 token
        cls_id = self.trx_dict[col]["cls_token_id"]
        sep_id = self.trx_dict[col]["sep_token_id"]
        input_ids = [cls_id] + padded_seq + [sep_id]  # [CLS, PAD, ..., tokens, SEP]

        # 4. [CLS] 始终在位置 0
        cls_pos = 0

        return input_ids, cls_pos

    def _process_seq(self, seq):
        result = {}

        # 处理时间戳
        raw_timestamps = seq[self.time_col]
        timestamps = self._parse_timestamps(raw_timestamps)

        # 处理事件序列
        event_type_col = self.cols_used[0]
        raw_event_types = seq[event_type_col]
        if isinstance(raw_event_types, str):
            try:
                event_types = ast.literal_eval(raw_event_types)
            except Exception as e:
                raise ValueError(f"Cannot parse event_types: {raw_event_types}") from e
        else:
            event_types = raw_event_types

        # 对齐长度
        seq_len = min(len(timestamps), len(event_types))
        event_types = event_types[:seq_len]
        timestamps = timestamps[:seq_len]

        # 截断序列
        max_len_wo_special = self.max_seq_len - 2
        if len(timestamps) > max_len_wo_special:
            truncated_event_types = event_types[-max_len_wo_special:]
            truncated_timestamps = timestamps[-max_len_wo_special:]
        else:
            truncated_event_types = event_types
            truncated_timestamps = timestamps

        # 计算时间特征
        age = self._compute_age_in_days(truncated_timestamps)
        interval = self._compute_interval_in_days(truncated_timestamps)
        same_event_interval = self._compute_same_event_interval_in_days(
            truncated_timestamps, truncated_event_types
        )

        # 填充时间特征：[CLS] 和 [SEP] 位置设为 0
        result['age'] = torch.tensor(
            self._pad_time_sequence(age.tolist()),
            dtype=torch.float, device=self.device
        )
        result['interval'] = torch.tensor(
            self._pad_time_sequence(interval.tolist()),
            dtype=torch.float, device=self.device
        )
        result['same_event_interval'] = torch.tensor(
            self._pad_time_sequence(same_event_interval.tolist()),
            dtype=torch.float, device=self.device
        )

        # 处理事件序列：改动点：接收 [CLS] 位置并存储
        for k in self.cols_used:
            input_ids, cls_pos = self._add_special_token(truncated_event_types, k)  # 接收 cls_pos
            result[k] = torch.tensor(input_ids, dtype=torch.long, device=self.device)
            result[f'{k}_mask'] = torch.tensor([1 if token != self.pad_token_id else 0 for token in input_ids],
                                              dtype=torch.long, device=self.device)
            result[f'{k}_cls_pos'] = torch.tensor(cls_pos, dtype=torch.long, device=self.device)  # 存储 [CLS] 位置

        for k in self.target_used:
            result[k] = torch.tensor(seq[k], dtype=torch.long, device=self.device)

        return result

if __name__ == '__main__':
    # 测试代码（验证 [CLS] 位置是否正确）
    trx_target_data = pd.DataFrame({
        "event_subtype": [
            [101, 205, 101, 302, 205],
            [401, 205, 401],
            [101, 101]
        ],
        "event_time": [
            [1634136341, 1638718103, 1663377587, 1667825293],
            [1669488361, 1671103076, 1671461313],
            [1609660447, 1610090831]
        ],
        "bcard_target":[1, 0, 0]  # 改为单标签（分类任务常用）
    })

    trx_dict = {
        "event_subtype": {
            "vocab_size": 500,
            "cls_token_id": 1,
            "sep_token_id": 2,
            "mask_token_id": 3
        }
    }

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    dataset = MBDDataset(
        trx_target_data=trx_target_data,
        trx_dict=trx_dict,
        max_seq_len=10,
        device=device,
        cols_used=['event_subtype'],
        target_used=['bcard_target'],
        time_col='event_time',
        time_unit='hour',
    )

    sample = dataset[0]
    print("input_ids:", sample['event_subtype'])  # 输出：[0, 1, 101, 205, 101, 2]（左补1个pad，[CLS]在索引1）
    print("[CLS]位置:", sample['event_subtype_cls_pos'])  # 输出：1（正确）
    print("attention_mask:", sample['event_subtype_mask'])